#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
通用内存监控工具
实时监控系统内存使用情况，并提供预警和自动清理功能
"""

import time
import argparse
import signal
import sys
import psutil
import gc
import logging
from datetime import datetime
from typing import Dict, Optional, Callable
from .utils import safe_print

logger = logging.getLogger(__name__)

class MemoryMonitor:
    """通用内存监控器"""
    
    def __init__(self, check_interval=30, warning_threshold=75, critical_threshold=85, 
                 cleanup_callback: Optional[Callable] = None):
        """
        初始化内存监控器
        
        Args:
            check_interval (int): 检查间隔（秒）
            warning_threshold (float): 警告阈值（百分比）
            critical_threshold (float): 严重阈值（百分比）
            cleanup_callback (Callable): 自定义清理回调函数
        """
        self.check_interval = check_interval
        self.warning_threshold = warning_threshold
        self.critical_threshold = critical_threshold
        self.cleanup_callback = cleanup_callback
        self.is_running = False
        self.total_warnings = 0
        self.total_cleanups = 0
        
        # 注册信号处理器
        signal.signal(signal.SIGINT, self._signal_handler)
        signal.signal(signal.SIGTERM, self._signal_handler)
    
    def _signal_handler(self, signum, frame):
        """信号处理器"""
        safe_print(f"\n📊 监控统计: 警告次数={self.total_warnings}, 清理次数={self.total_cleanups}")
        safe_print("🛑 内存监控器已停止")
        self.is_running = False
        sys.exit(0)
    
    def get_memory_usage(self) -> Optional[Dict]:
        """获取内存使用情况"""
        try:
            # 系统内存信息
            system_memory = psutil.virtual_memory()
            
            # 当前进程内存信息
            current_process = psutil.Process()
            process_memory = current_process.memory_info()
            
            return {
                'system': {
                    'total': system_memory.total / (1024**3),  # GB
                    'used': system_memory.used / (1024**3),   # GB
                    'available': system_memory.available / (1024**3),  # GB
                    'percent': system_memory.percent
                },
                'process': {
                    'rss': process_memory.rss / (1024**2),    # MB
                    'vms': process_memory.vms / (1024**2),    # MB
                    'percent': current_process.memory_percent()
                }
            }
        except Exception as e:
            logger.error(f"获取内存使用情况失败: {e}")
            return None
    
    def default_cleanup(self, reason: str = "内存清理"):
        """默认内存清理函数"""
        try:
            safe_print(f"🧹 执行{reason}...")
            
            # 强制垃圾回收
            collected = gc.collect()
            safe_print(f"   垃圾回收: 清理了 {collected} 个对象")
            
            # 清理未使用的内存
            if hasattr(gc, 'set_threshold'):
                # 调整垃圾回收阈值，更积极地回收内存
                gc.set_threshold(700, 10, 10)
            
            return True
        except Exception as e:
            logger.error(f"内存清理失败: {e}")
            return False
    
    def cleanup_resources(self, reason: str = "内存清理"):
        """清理资源"""
        if self.cleanup_callback:
            try:
                return self.cleanup_callback(reason)
            except Exception as e:
                logger.error(f"自定义清理回调执行失败: {e}")
                return self.default_cleanup(reason)
        else:
            return self.default_cleanup(reason)
    
    def check_memory_status(self):
        """检查内存状态"""
        memory_info = self.get_memory_usage()
        
        if memory_info is None:
            safe_print("❌ 无法获取内存信息")
            return None
        
        # 获取当前时间
        current_time = datetime.now().strftime("%H:%M:%S")
        
        # 系统内存使用情况
        system_percent = memory_info['system']['percent']
        system_used = memory_info['system']['used']
        system_total = memory_info['system']['total']
        
        # 进程内存使用情况
        process_rss = memory_info['process']['rss']
        process_percent = memory_info['process']['percent']
        
        # 根据内存使用率确定显示样式
        if system_percent >= self.critical_threshold:
            status_icon = "🔴"
            status_text = "严重"
            self.total_warnings += 1
            
            # 执行自动清理
            safe_print(f"\n{status_icon} [{current_time}] {status_text} - 内存使用率过高!")
            safe_print(f"   系统内存: {system_percent:.1f}% ({system_used:.1f}GB/{system_total:.1f}GB)")
            safe_print(f"   进程内存: {process_percent:.1f}% ({process_rss:.1f}MB)")
            
            if self.cleanup_resources("自动内存清理"):
                self.total_cleanups += 1
                
                # 清理后重新检查
                time.sleep(2)
                new_memory_info = self.get_memory_usage()
                if new_memory_info:
                    new_system_percent = new_memory_info['system']['percent']
                    new_process_rss = new_memory_info['process']['rss']
                    memory_freed = process_rss - new_process_rss
                    safe_print(f"✅ 清理完成 - 系统内存: {new_system_percent:.1f}%, "
                              f"释放进程内存: {memory_freed:.1f}MB\n")
            
        elif system_percent >= self.warning_threshold:
            status_icon = "🟡"
            status_text = "警告"
            self.total_warnings += 1
            safe_print(f"{status_icon} [{current_time}] {status_text} - 系统: {system_percent:.1f}%, "
                      f"进程: {process_rss:.1f}MB")
        else:
            status_icon = "🟢"
            status_text = "正常"
            safe_print(f"{status_icon} [{current_time}] {status_text} - 系统: {system_percent:.1f}%, "
                      f"进程: {process_rss:.1f}MB")
        
        return memory_info
    
    def start_monitoring(self):
        """开始监控"""
        safe_print("🔍 通用内存监控器")
        safe_print("=" * 50)
        safe_print(f"📊 监控参数:")
        safe_print(f"   检查间隔: {self.check_interval}秒")
        safe_print(f"   警告阈值: {self.warning_threshold}%")
        safe_print(f"   严重阈值: {self.critical_threshold}%")
        safe_print(f"   自动清理: {'自定义' if self.cleanup_callback else '默认'}")
        safe_print("=" * 50)
        safe_print("按 Ctrl+C 停止监控\n")
        
        self.is_running = True
        
        while self.is_running:
            try:
                self.check_memory_status()
                time.sleep(self.check_interval)
            except KeyboardInterrupt:
                break
            except Exception as e:
                safe_print(f"❌ 监控过程中出错: {e}")
                logger.error(f"监控过程中出错: {e}")
                time.sleep(self.check_interval)
        
        safe_print(f"\n📊 监控统计: 警告次数={self.total_warnings}, 清理次数={self.total_cleanups}")
        safe_print("🛑 内存监控器已停止")
    
    def stop_monitoring(self):
        """停止监控"""
        self.is_running = False
    
    def get_current_status(self) -> Dict:
        """获取当前内存状态（用于API调用）"""
        memory_info = self.get_memory_usage()
        if not memory_info:
            return {'error': '无法获取内存信息'}
        
        system_percent = memory_info['system']['percent']
        
        if system_percent >= self.critical_threshold:
            status = 'critical'
            level = '严重'
        elif system_percent >= self.warning_threshold:
            status = 'warning'
            level = '警告'
        else:
            status = 'normal'
            level = '正常'
        
        return {
            'status': status,
            'level': level,
            'system_percent': system_percent,
            'system_used_gb': memory_info['system']['used'],
            'system_total_gb': memory_info['system']['total'],
            'process_rss_mb': memory_info['process']['rss'],
            'process_percent': memory_info['process']['percent'],
            'timestamp': datetime.now().isoformat(),
            'statistics': {
                'total_warnings': self.total_warnings,
                'total_cleanups': self.total_cleanups
            }
        }

def parse_arguments():
    """解析命令行参数"""
    parser = argparse.ArgumentParser(description="通用内存监控器")
    
    parser.add_argument(
        '--interval', '-i',
        type=int,
        default=30,
        help='检查间隔（秒），默认30秒'
    )
    
    parser.add_argument(
        '--warning',
        type=float,
        default=75.0,
        help='警告阈值（百分比），默认75%%'
    )
    
    parser.add_argument(
        '--critical',
        type=float,
        default=85.0,
        help='严重阈值（百分比），默认85%%'
    )
    
    parser.add_argument(
        '--once',
        action='store_true',
        help='只检查一次内存状态，不持续监控'
    )
    
    parser.add_argument(
        '--no-cleanup',
        action='store_true',
        help='禁用自动内存清理'
    )
    
    return parser.parse_args()

def main():
    """主函数"""
    args = parse_arguments()
    
    # 验证参数
    if args.warning >= args.critical:
        safe_print("❌ 错误: 警告阈值必须小于严重阈值")
        sys.exit(1)
    
    if args.interval < 1:
        safe_print("❌ 错误: 检查间隔必须大于0秒")
        sys.exit(1)
    
    # 创建监控器
    cleanup_callback = None if args.no_cleanup else None
    monitor = MemoryMonitor(
        check_interval=args.interval,
        warning_threshold=args.warning,
        critical_threshold=args.critical,
        cleanup_callback=cleanup_callback
    )
    
    if args.once:
        # 只检查一次
        safe_print("🔍 检查当前内存状态...")
        status = monitor.get_current_status()
        safe_print(f"状态: {status.get('level', 'N/A')}")
        safe_print(f"系统内存: {status.get('system_percent', 0):.1f}%")
        safe_print(f"进程内存: {status.get('process_rss_mb', 0):.1f}MB")
    else:
        # 持续监控
        monitor.start_monitoring()

if __name__ == "__main__":
    main()